"""
嵌入模型服务
支持豆包API调用模式
"""
from typing import List, Union
import requests
from loguru import logger

from app.core.config import settings


class DoubaoEmbeddingService:
    """豆包嵌入服务"""

    def __init__(self, api_key: str, model_name: str, base_url: str):
        self.api_key = api_key
        self.model_name = model_name
        self.base_url = base_url.rstrip('/')

        logger.info(f"初始化豆包嵌入服务: {model_name} @ {base_url}")

    def encode(self, texts: Union[str, List[str]]) -> Union[List[float], List[List[float]]]:
        """编码文本为向量"""
        try:
            # 确保输入是列表格式
            if isinstance(texts, str):
                input_texts = [texts]
                return_single = True
            else:
                input_texts = texts
                return_single = False

            # 构建豆包API请求格式
            headers = {
                "Authorization": f"Bearer {self.api_key}",
                "Content-Type": "application/json"
            }

            # 对于批量请求，分别处理每个文本（豆包API每次请求返回一个向量）
            embeddings = []

            for text in input_texts:
                # 使用标准文本嵌入API
                data = {
                    "model": self.model_name,
                    "encoding_format": "float",
                    "input": text
                }

                # 发送请求到豆包文本嵌入API
                response = requests.post(
                    f"{self.base_url}/embeddings",
                    headers=headers,
                    json=data,
                    timeout=30
                )
                response.raise_for_status()

                # 解析豆包API响应
                result = response.json()

                # 豆包API返回嵌入向量列表
                if "data" in result and len(result["data"]) > 0 and "embedding" in result["data"][0]:
                    embedding = result["data"][0]["embedding"]
                    embeddings.append(embedding)
                else:
                    raise ValueError(f"豆包API响应格式错误: {result}")

            if return_single:
                return embeddings[0]
            else:
                return embeddings

        except Exception as e:
            logger.error(f"豆包嵌入编码失败: {e}")
            raise

    def is_available(self) -> bool:
        """检查服务是否可用"""
        try:
            # 简单的健康检查 - 发送一个测试请求
            headers = {"Authorization": f"Bearer {self.api_key}"}
            test_data = {
                "model": self.model_name,
                "encoding_format": "float",
                "input": "test"
            }

            response = requests.post(
                f"{self.base_url}/embeddings",
                headers=headers,
                json=test_data,
                timeout=10
            )
            return response.status_code == 200
        except:
            return False


def create_embedding_service() -> DoubaoEmbeddingService:
    """创建豆包嵌入服务实例"""

    if not settings.embedding_api_key:
        raise ValueError("需要设置 EMBEDDING_API_KEY")

    return DoubaoEmbeddingService(
        api_key=settings.embedding_api_key,
        model_name=settings.embedding_model,
        base_url=settings.embedding_api_base_url
    )
